import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk
import serial.tools.list_ports
import subprocess
import threading
import queue
import re
import tempfile
import os
import sys


class ESP32Flasher(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("ESP32 Flasher")
        self.geometry("500x350")  # 设置默认窗口大小

        # 已完成烧录次数
        self.flash_count = 0

        self.create_widgets()
        self.queue = queue.Queue()
        self.check_queue()

        # 绑定窗口关闭事件
        self.protocol("WM_DELETE_WINDOW", self.on_closing)

        # 初始化进程变量
        self.process = None

        # 开始动态刷新COM口
        self.refresh_serial_ports()

    def create_widgets(self):
        # 使用grid布局管理器
        self.grid_columnconfigure(1, weight=1)

        # File selection
        self.file_label = tk.Label(self, text="选择固件文件:")
        self.file_label.grid(row=0, column=0, padx=5, pady=5, sticky=tk.W)

        self.file_entry = tk.Entry(self, width=40)
        self.file_entry.grid(row=1, column=1, padx=5, pady=5, sticky=tk.W)

        self.browse_button = tk.Button(self, text="浏览", command=self.browse_file)
        self.browse_button.grid(row=1, column=2, padx=5, pady=5, sticky=tk.W)

        # Port selection
        self.port_label = tk.Label(self, text="选择串口:")
        self.port_label.grid(row=2, column=0, padx=5, pady=5, sticky=tk.W)

        self.port_combobox = tk.StringVar()
        self.port_combobox.set("选择串口")  # 设置默认值

        self.port_menu = tk.OptionMenu(self, self.port_combobox, "未检测到串口")
        self.port_menu.grid(row=2, column=1, padx=5, pady=5, sticky=tk.W)

        # MAC address input
        self.mac_label = tk.Label(self, text="MAC地址 (6字节16进制):")
        self.mac_label.grid(row=3, column=0, padx=5, pady=5, sticky=tk.W)

        self.mac_entry = tk.Entry(self, width=40)
        self.mac_entry.grid(row=3, column=1, padx=5, pady=5, sticky=tk.W)

        # Progress bar
        self.progress = ttk.Progressbar(self, orient="horizontal", length=300, mode="determinate")
        self.progress.grid(row=4, column=0, columnspan=3, padx=5, pady=10, sticky=tk.W + tk.E)

        # Flash button
        self.flash_button = tk.Button(self, text="烧录固件", command=self.flash_firmware)
        self.flash_button.grid(row=5, column=0, columnspan=3, padx=5, pady=20, sticky=tk.W + tk.E)

        # Flash count label
        self.flash_count_label = tk.Label(self, text=f"已烧录次数: {self.flash_count}")
        self.flash_count_label.grid(row=6, column=0, columnspan=3, padx=5, pady=5, sticky=tk.W + tk.E)

    def browse_file(self):
        file_path = filedialog.askopenfilename(filetypes=[("Binary files", "*.bin")])
        if file_path:
            self.file_entry.delete(0, tk.END)
            self.file_entry.insert(0, file_path)

    def get_serial_ports(self):
        ports = serial.tools.list_ports.comports()
        return [port.device for port in ports]

    def refresh_serial_ports(self):
        current_selection = self.port_combobox.get()
        current_ports = self.get_serial_ports()
        menu = self.port_menu["menu"]
        menu.delete(0, "end")
        
        if current_ports:
            for port in current_ports:
                menu.add_command(label=port, command=lambda p=port: self.port_combobox.set(p))
            # 如果当前选择的端口仍然存在，保持选中状态，否则选第一个
            if current_selection in current_ports:
                self.port_combobox.set(current_selection)
            else:
                self.port_combobox.set(current_ports[0])
        else:
            menu.add_command(label="未检测到串口", command=lambda: self.port_combobox.set("未检测到串口"))
            self.port_combobox.set("未检测到串口")

        # 继续每秒刷新
        self.after(1000, self.refresh_serial_ports)

    def validate_mac(self, mac):
        # 检查MAC地址是否为6个字节的16进制数
        return bool(re.fullmatch(r'[0-9A-Fa-f]{12}', mac))

    def increment_mac(self, mac):
        # 将MAC地址转换为整数，自加1，再转换回16进制字符串
        mac_int = int(mac, 16) + 1
        mac_hex = f'{mac_int:012X}'
        return mac_hex

    def flash_firmware(self):
        file_path = self.file_entry.get()
        port = self.port_combobox.get()
        raw_mac = self.mac_entry.get().strip()
        mac_address = raw_mac.replace(":", "").replace("-", "")
        
        if not file_path or not port or port == "未检测到串口" or port == "选择串口":
            messagebox.showerror("错误", "请选择固件文件和串口")
            return

        if mac_address and not self.validate_mac(mac_address):
            messagebox.showerror("错误", "MAC地址格式错误，请输入6字节16进制数")
            return

        # Disable the flash button to prevent multiple clicks
        self.flash_button.config(state=tk.DISABLED)

        # Start the flashing process in a separate thread
        threading.Thread(target=self.run_esptool, args=(port, file_path, mac_address)).start()

    def run_esptool(self, port, file_path, mac_address):
        try:
            self.progress["value"] = 0
            self.progress["maximum"] = 100
            # 获取打包后的资源路径
            base_path = getattr(sys, '_MEIPASS', os.path.abspath("."))
            esptool_path = os.path.join(base_path, "resources", "esptool.exe")
            
            # 只有当MAC地址非空时才烧录MAC
            if mac_address:
                print(mac_address)

                # 创建临时文件（安全方式）
                with tempfile.NamedTemporaryFile(mode='wb', suffix='.bin', delete=False) as mac_file:
                    mac_bytes = bytes.fromhex(mac_address)
                    mac_file.write(mac_bytes)
                    mac_path = mac_file.name

                # 擦除区域
                erase_cmd = [esptool_path, "--port", port, "erase_region", "0x3FE000", "0x1000"]
                erase_process = subprocess.Popen(
                    erase_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True
                )
                erase_process.wait()
                
                if erase_process.returncode != 0:
                    error = erase_process.stderr.read()
                    self.queue.put(f"ERROR: 擦除失败 - {error}")
                    return

                # 烧录MAC地址
                mac_cmd = [esptool_path, "--port", port, "write_flash", "0x3FE000", mac_path]
                mac_process = subprocess.Popen(
                    mac_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True
                )
                mac_process.wait()
                
                if mac_process.returncode != 0:
                    error = mac_process.stderr.read()
                    self.queue.put(f"ERROR: MAC烧录失败 - {error}")
                    return
            else :
                print("No MAC Address")

            self.process = subprocess.Popen(
                [esptool_path, "--chip", "esp32", "--port", port, "-b", "921600", "--before=default_reset", "--after=hard_reset",
                 "write_flash", "--flash_mode", "dio", "--flash_freq", "80m", "--flash_size", "4MB", "0", file_path],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            for line in iter(self.process.stdout.readline, ''):
                if line:
                    self.queue.put(line)
            self.process.stdout.close()
            self.process.wait()

            if self.process.returncode == 0:
                self.queue.put("SUCCESS")
            else:
                stderr_output = self.process.stderr.read()
                self.queue.put(f"ERROR: {stderr_output}")

        except Exception as e:
            self.queue.put(f"ERROR: {e}")
        finally:
            # Re-enable the flash button
            self.queue.put("DONE")

    def check_queue(self):
        try:
            while True:
                line = self.queue.get_nowait()
                if line == "SUCCESS":
                    # 只有当存在MAC地址时才自增
                    current_mac = self.mac_entry.get().strip()
                    if current_mac:
                        clean_mac = current_mac.replace(":", "").replace("-", "")
                        new_mac = self.increment_mac(clean_mac)
                        self.mac_entry.delete(0, tk.END)
                        self.mac_entry.insert(0, new_mac)
                    
                    self.flash_count += 1
                    self.flash_count_label.config(text=f"已烧录次数: {self.flash_count}")
                    messagebox.showinfo("成功", "固件烧录成功")

                elif line.startswith("ERROR"):
                    self.show_message("错误", line)

                elif line == "DONE":
                    self.flash_button.config(state=tk.NORMAL)

                else:
                    self.update_progress(line)

        except queue.Empty:
            pass
        self.after(100, self.check_queue)

    def update_progress(self, line):
        # Parse the progress from the esptool output
        if "Writing" in line and "%" in line:
            print(line)  # 打印调试信息
            progress_str = line.split("(")[-1].split("%")[0].strip()
            try:
                progress = float(progress_str)
                self.progress["value"] = progress
            except ValueError:
                pass

    def show_message(self, title, message):
        messagebox.showinfo(title, message)

    def on_closing(self):
        # 如果有正在运行的进程，终止它
        if self.process and self.process.poll() is None:
            try:
                self.process.terminate()
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()
        self.destroy()  # 关闭窗口

if __name__ == "__main__":
    app = ESP32Flasher()
    app.mainloop()
